import torch
import numpy as np
from environment import MazeEnvironment
from config import *
from graph_utils import create_distance_maps, create_direction_maps, grid_to_graph_features


def generate_dataset(num_samples=NUM_SAMPLES, seed=RANDOM_SEED):
    """生成训练、验证和测试数据集，优先采样智能体移动的样本，包含增强的状态表示"""
    # 使用固定种子创建环境
    env = MazeEnvironment(seed=seed)

    # 保存初始布局用于后续使用
    initial_layout = env.get_state().copy()
    np.save('initial_layout.npy', initial_layout)

    # 创建静态掩码 - 标记不应变化的元素
    static_mask = np.zeros_like(initial_layout, dtype=bool)
    static_mask[initial_layout == OBSTACLE] = True  # 障碍物不变
    static_mask[initial_layout == START] = True  # 起点不变
    static_mask[initial_layout == END] = True  # 终点不变
    static_mask[initial_layout == COIN] = True  # 金币位置不变
    np.save('static_mask.npy', static_mask)

    states = []
    actions = []
    next_states = []
    agent_positions = []  # 绝对位置
    agent_movements = []  # 相对移动 (dx, dy)

    # 记录已经见过的状态转换
    seen_transitions = set()

    # 跟踪移动样本的比例
    movement_samples = 0
    total_samples = 0

    # 生成样本，直到达到所需数量
    while len(states) < num_samples:
        # 定期重置智能体位置，但保持环境布局不变
        if len(states) % 50 == 0:
            env.reset_without_changing_layout()

        # 获取当前状态和智能体位置
        current_state = env.get_state().copy()
        current_agent_pos = env.get_agent_position(current_state)

        if current_agent_pos is None:
            env.reset_without_changing_layout()
            continue

        # 优先选择导致智能体移动的动作
        valid_actions = []
        move_actions = []

        for action in range(4):  # 只考虑UP, DOWN, LEFT, RIGHT
            next_state = env.get_next_state_after_action(action)
            next_agent_pos = env.get_agent_position(next_state)

            if next_agent_pos and next_agent_pos != current_agent_pos:
                move_actions.append(action)

            if next_agent_pos:
                valid_actions.append(action)

        # 如果有移动动作，大部分时间选择移动动作
        if move_actions and (movement_samples < AGENT_MOVEMENT_RATIO * total_samples or np.random.random() < 0.9):
            action = np.random.choice(move_actions)
            is_movement = True
        elif valid_actions:
            action = np.random.choice(valid_actions)
            is_movement = (action in move_actions)
        else:
            action = np.random.randint(0, 4)
            is_movement = False

        # 获取下一个状态
        next_state = env.get_next_state_after_action(action)
        next_agent_pos = env.get_agent_position(next_state)

        # 只有当能找到智能体位置时才添加样本
        if current_agent_pos and next_agent_pos:
            # 计算相对移动
            dx = next_agent_pos[0] - current_agent_pos[0]
            dy = next_agent_pos[1] - current_agent_pos[1]

            # 创建增强的状态表示
            current_input = create_enhanced_state_input(current_state)
            next_input = create_enhanced_state_input(next_state)

            # 添加样本
            states.append(current_input)
            actions.append(action)
            next_states.append(next_input.flatten())
            agent_positions.append([next_agent_pos[0], next_agent_pos[1]])
            agent_movements.append([dx, dy, next_agent_pos[0], next_agent_pos[1]])

            # 更新计数器
            total_samples += 1
            if is_movement:
                movement_samples += 1

            # 执行动作
            env.step(action)
        else:
            # 如果找不到智能体，重置环境
            env.reset_without_changing_layout()

    # 确保样本数量正确
    states = states[:num_samples]
    actions = actions[:num_samples]
    next_states = next_states[:num_samples]
    agent_positions = agent_positions[:num_samples]
    agent_movements = agent_movements[:num_samples]

    # 输出移动样本比例
    movement_ratio = movement_samples / total_samples
    print(f"生成的数据集中包含智能体移动的样本比例: {movement_ratio:.2f}")

    # 转换为numpy数组
    states = np.array(states)
    actions = np.array(actions)
    next_states = np.array(next_states)
    agent_positions = np.array(agent_positions)
    agent_movements = np.array(agent_movements)

    # 随机打乱数据，确保数据分布均匀
    rng = np.random.RandomState(seed)
    indices = rng.permutation(num_samples)
    states = states[indices]
    actions = actions[indices]
    next_states = next_states[indices]
    agent_positions = agent_positions[indices]
    agent_movements = agent_movements[indices]

    # 分割为训练、验证和测试集
    num_train = int(num_samples * TRAIN_SPLIT)
    num_val = int(num_samples * VAL_SPLIT)

    train_indices = np.arange(num_train)
    val_indices = np.arange(num_train, num_train + num_val)
    test_indices = np.arange(num_train + num_val, num_samples)

    train_data = (states[train_indices], actions[train_indices], next_states[train_indices],
                  agent_positions[train_indices], agent_movements[train_indices])
    val_data = (states[val_indices], actions[val_indices], next_states[val_indices],
                agent_positions[val_indices], agent_movements[val_indices])
    test_data = (states[test_indices], actions[test_indices], next_states[test_indices],
                 agent_positions[test_indices], agent_movements[test_indices])

    return train_data, val_data, test_data


def create_enhanced_state_input(grid_state):
    """创建增强版的状态输入，包含距离和方向信息"""
    # 计算需要的特征总数
    total_features = CELL_FEATURES

    # 计算距离图
    dist_maps = []
    if USE_DISTANCE_MAPS:
        dist_maps = create_distance_maps(grid_state)
        total_features += len(dist_maps)

    # 计算方向图
    dir_maps = []
    if USE_DIRECTION_MAPS:
        dir_maps = create_direction_maps(grid_state)
        total_features += len(dir_maps)

    # 创建具有足够通道的增强状态表示
    enhanced_input = np.zeros((ENV_SIZE, ENV_SIZE, total_features))

    # 首先为每个单元格分配基本特征
    for i in range(ENV_SIZE):
        for j in range(ENV_SIZE):
            cell_type = grid_state[i, j]
            enhanced_input[i, j, cell_type] = 1

    # 添加距离图
    offset = CELL_FEATURES
    for idx, dist_map in enumerate(dist_maps):
        for i in range(ENV_SIZE):
            for j in range(ENV_SIZE):
                enhanced_input[i, j, offset + idx] = dist_map[i, j]

    # 添加方向图
    offset += len(dist_maps)
    for idx, dir_map in enumerate(dir_maps):
        for i in range(ENV_SIZE):
            for j in range(ENV_SIZE):
                enhanced_input[i, j, offset + idx] = dir_map[i, j]

    return enhanced_input.flatten()


def prepare_batch(states, actions, next_states, agent_positions=None, agent_movements=None):
    """准备批次数据"""
    states_tensor = torch.FloatTensor(states)
    actions_tensor = torch.LongTensor(actions)
    next_states_tensor = torch.FloatTensor(next_states)

    if agent_positions is not None and agent_movements is not None:
        agent_positions_tensor = torch.FloatTensor(agent_positions)
        agent_movements_tensor = torch.FloatTensor(agent_movements)
        return states_tensor, actions_tensor, next_states_tensor, agent_positions_tensor, agent_movements_tensor
    elif agent_positions is not None:
        agent_positions_tensor = torch.FloatTensor(agent_positions)
        return states_tensor, actions_tensor, next_states_tensor, agent_positions_tensor

    return states_tensor, actions_tensor, next_states_tensor


def calculate_accuracy(outputs, targets):
    """计算准确率"""
    predicted = torch.argmax(outputs, dim=1)
    correct = (predicted == targets).float()
    return torch.mean(correct).item()


def extract_agent_positions(states, env_size=ENV_SIZE):
    """从一批状态中提取agent位置"""
    batch_size = states.shape[0]
    agent_positions = []

    for i in range(batch_size):
        state = states[i].reshape(env_size, env_size)
        agent_pos = np.where(state == AGENT)
        if len(agent_pos[0]) > 0:
            agent_positions.append((agent_pos[0][0], agent_pos[1][0]))
        else:
            agent_positions.append((-1, -1))  # 表示找不到agent

    return agent_positions


def calculate_agent_position_accuracy(pred_positions, true_positions):
    """计算智能体绝对位置预测的准确率"""
    correct = 0
    total = 0

    for pred, true in zip(pred_positions, true_positions):
        if true[0] == -1 or true[1] == -1:  # 如果找不到真实agent位置
            continue

        if pred[0] == true[0] and pred[1] == true[1]:
            correct += 1
        total += 1

    return correct / max(1, total)


def calculate_agent_movement_accuracy(pred_movements, true_movements):
    """计算智能体相对移动预测的准确率"""
    correct = 0
    total = 0

    for pred, true in zip(pred_movements, true_movements):
        pred_dx, pred_dy = pred[0], pred[1]
        true_dx, true_dy = true[0], true[1]

        if pred_dx == true_dx and pred_dy == true_dy:
            correct += 1
        total += 1

    return correct / max(1, total)